"""
    多进程执行基因簇识别和计算基因密度脚本
"""


import os
import time
import multiprocessing
from multiprocessing import Pool    


# 读入imp每个物种信息文件夹及文件夹的路径，文件名是元组的第一个元素，路径是第二个元素
# 每个物种的信息包括，bed文件、基因数量与碱基熟练文件、基因的Pfam注释结果表格
files = os.listdir("../imp_data")
dir_path = [os.path.join("../imp_data",x).replace("\\","/") for x in files]
species_path = zip(files,dir_path)


# 调用植物基因组上基于密度识别基因簇的程序生成结果
def run_cluster(species_name, all_base_num, all_gene_num, bed_path, output_table_dir):
    cmd = (
        "python ./01.plant_gene_cluster_base_local_density.py " + 
        " -species " +  species_name + 
        " -base_num " + all_base_num + 
        " -gene_num " + all_gene_num + 
        " -bed " + bed_path +
        " -output_table " + output_table_dir 
    )
    os.system(cmd)


if __name__ == "__main__":
    start_time = time.time()

    task_list = []
    for species, path in species_path:
        counts_length_file = path + "/" + species + "_mRNA_countsLength.txt"
        with open(counts_length_file) as file:
            gene_base_num = file.read().strip().split("\t")
        all_base_num = gene_base_num[2]
        all_gene_num = gene_base_num[1]
        bed_path = path + "/" + species + ".bed"
        output_table_dir = "./01.cluster_result"
        os.makedirs(output_table_dir, exist_ok=True)
        task_tuple = (species, all_base_num, all_gene_num, bed_path, output_table_dir)
        task_list.append(task_tuple)

    cpu_count = multiprocessing.cpu_count() -1
    with Pool(20) as pool:
        pool.starmap(run_cluster, task_list[0:])
        
    end_time = time.time()
    total_time = end_time - start_time 
    
    print(
        f"\nThe number of tasks is {len(list(task_list))} :\n", 
        f"\t Time taken {round(total_time, 2)}s\n"
    )